# Copyright (c) 2020 Presto Labs Pte. Ltd.
# Author: jingyuan

from collections import defaultdict

from ..on_log import OnLog
from .base import ExportDataFrame, OrderInfoTracker


class CancelNumber(OrderInfoTracker, ExportDataFrame, OnLog):
  def __init__(self, cancel_limit=5):
    super().__init__()
    self.cancel_limit = cancel_limit
    self.cancel_count = defaultdict(int)
    self.errors = {}
    self.cancel_errors = {}
    self.ts_fields += ['last_cancel_timestamp']

  def on_cancel_submit(self, ts, symbol, order_id):
    self.cancel_count[order_id] += 1
    if self.cancel_count[order_id] > self.cancel_limit and order_id in self.order_info:
      self.errors[order_id] = {
          'cancel_count': self.cancel_count[order_id],
          'last_cancel_timestamp': ts,
          'order_id': order_id
      }
      self.errors[order_id].update(self.order_info[order_id])

  def on_cancel_error(self, ts, currency, order_id, *args):
    names = ['error_code', 'exchange_error_code', 'error_msg']
    if order_id not in self.cancel_errors:
      self.cancel_errors[order_id] = {n: '' for n in names}
    for i, name in enumerate(names):
      if str(args[i]) not in self.cancel_errors[order_id][name]:
        if len(self.cancel_errors[order_id][name]) > 0:
          self.cancel_errors[order_id][name] += f',{args[i]}'
        else:
          self.cancel_errors[order_id][name] += f'{args[i]}'

  def on_cancel_reject(self, *args, **kwargs):
    self.on_cancel_error(*args, **kwargs)

  def to_df(self):
    for oid in self.errors:
      if oid in self.cancel_errors:
        self.errors[oid].update(self.cancel_errors[oid])
    return super().to_df().sort_values('cancel_count', ascending=False)
